#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.1
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : ExcelExporter.py
# @Software: PyCharm
# @time: 2025/5/30 9:54
# @description : Excel导出工具类
"""

import os
import logging
import pandas as pd
import openpyxl
from openpyxl.styles import Font, Alignment, PatternFill, Border, Side
from openpyxl.utils import get_column_letter
from datetime import datetime
import re
from typing import Dict, List, Any, Union

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class ExcelExporter:
    """Excel导出类，负责将数据导出到Excel文件"""

    def __init__(self):
        # 定义样式
        self.styles = {
            'header_font': Font(bold=True, color="FFFFFF", size=12),
            'header_fill': PatternFill(start_color="4F81BD", end_color="4F81BD", fill_type="solid"),
            'header_alignment': Alignment(horizontal='center', vertical='center', wrap_text=True),
            'data_font': Font(name='Arial', size=10),
            'data_alignment': Alignment(horizontal='center', vertical='center'),
            'positive_fill': PatternFill(start_color="C6EFCE", end_color="C6EFCE", fill_type="solid"),
            'negative_fill': PatternFill(start_color="FFC7CE", end_color="FFC7CE", fill_type="solid"),
            'thin_border': Border(left=Side(style='thin'),
                                  right=Side(style='thin'),
                                  top=Side(style='thin'),
                                  bottom=Side(style='thin')),
            'date_format': 'yyyy-mm-dd',
            'number_format': '#,##0.00',
            'percent_format': '0.00%'
        }

    def clean_target_date(self, target_date: Union[str, datetime, pd.Timestamp]) -> str:
        """
        清理目标日期，返回合法的文件名格式的日期字符串
        :param target_date: 目标日期，可能是字符串、datetime或pandas Timestamp
        :return: 格式化的日期字符串
        """
        if isinstance(target_date, (pd.Timestamp, datetime)):
            return target_date.strftime("%Y%m%d")
        elif isinstance(target_date, str):
            clean_date = target_date.split()[0]
            return clean_date.replace(":", "_").replace("/", "-").replace("\\", "_")
        else:
            return str(target_date).split()[0].replace(":", "_").replace("/", "-").replace("\\", "_")

    def export_signals_report(self, signals: Dict[str, List[Dict[str, Any]]],
                              target_date: Union[str, datetime, pd.Timestamp],
                              output_dir: str = "reports") -> str:
        """
        生成信号Excel报告
        :param signals: 信号数据，包含'sz_signals'和'sh_signals'
        :param target_date: 目标日期
        :param output_dir: 输出目录
        :return: Excel文件路径
        """
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)

        # 创建时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

        # 处理目标日期
        clean_target_date_str = self.clean_target_date(target_date=target_date)

        # 创建文件名 - 只使用安全的字符
        filename = f"股票信号报告_{clean_target_date_str}_{timestamp}.xlsx"
        filepath = os.path.join(output_dir, filename)

        # 创建Excel工作簿
        wb = openpyxl.Workbook()

        # 处理深圳市场信号
        if signals.get('sz_signals'):
            sz_df = self.process_signals(signals['sz_signals'], target_date)
            if not sz_df.empty:
                ws_sz = wb.create_sheet(title="深圳市场信号")
                self.add_data_to_sheet(ws_sz, sz_df, "深圳市场信号")

        # 处理上海市场信号
        if signals.get('sh_signals'):
            sh_df = self.process_signals(signals['sh_signals'], target_date)
            if not sh_df.empty:
                ws_sh = wb.create_sheet(title="上海市场信号")
                self.add_data_to_sheet(ws_sh, sh_df, "上海市场信号")

        # 删除默认创建的工作表
        if 'Sheet' in wb.sheetnames:
            del wb['Sheet']

        # 保存文件
        try:
            wb.save(filepath)
            logging.info(f"Excel报告已生成: {filepath}")
        except Exception as e:
            logging.error(f"保存Excel报告失败: {e}")
            raise

        return filepath

    def process_signals(self, signals_list: List[Dict[str, Any]], target_date: str) -> pd.DataFrame:
        """
        处理信号数据并创建DataFrame
        :param signals_list: 信号列表
        :param target_date: 目标日期
        :return: 处理后的DataFrame
        """
        if not signals_list:
            return pd.DataFrame()

        # 创建合并信号字典
        combined_signals = {}
        for signal in signals_list:
            code = signal['stock_code']
            if code not in combined_signals:
                combined_signals[code] = {
                    '日期': signal.get('date', '未知'),
                    '股票代码': code,
                    '股票名称': signal.get('stock_name', '未知'),
                    '评级': signal.get('rating', 'N/A'),
                    '行业': signal.get('industry', '未知'),
                    '收盘价': signal.get('close_price', 0),
                    '成交量(万手)': signal.get('volume', 0),
                    '涨跌幅(%)': signal.get('change_pct', 0),
                    '信号列表': [],
                    '信号描述': []
                }

            # 添加信号信息
            signal_name = signal.get('signal_name', '未知信号')
            signal_desc = signal.get('description', '无描述')
            combined_signals[code]['信号列表'].append(signal_name)
            combined_signals[code]['信号描述'].append(signal_desc)

        # 转换为列表格式
        data_list = []
        for code, info in combined_signals.items():
            # 合并信号信息
            signals_str = "; ".join(info['信号列表'])
            desc_str = "; ".join(info['信号描述'])
            # 将涨跌幅(%)的值除以100并保留两位小数
            change_rate = round(info['涨跌幅(%)'] / 100, 4)

            data_list.append({
                '日期': info['日期'],
                '股票代码': info['股票代码'],
                '股票名称': info['股票名称'],
                '评级': info['评级'],
                '行业': info['行业'],
                '收盘价': info['收盘价'],
                '成交量(万手)': info['成交量(万手)'],
                '涨跌幅(%)': change_rate,
                '信号列表': signals_str,
                '信号描述': desc_str
            })

        # 创建DataFrame
        df = pd.DataFrame(data_list)

        # 按行业和股票代码排序
        if not df.empty:
            df.sort_values(by=['行业', '股票代码'], inplace=True)

        return df

    def add_data_to_sheet(self, worksheet: openpyxl.worksheet.worksheet.Worksheet,
                         df: pd.DataFrame, title: str) -> None:
        """
        将DataFrame数据添加到工作表并应用样式
        :param worksheet: 工作表对象
        :param df: 包含数据的DataFrame
        :param title: 工作表标题
        """
        # 添加标题行
        headers = list(df.columns)
        worksheet.append(headers)

        # 应用标题样式
        for col_idx, header in enumerate(headers, 1):
            cell = worksheet.cell(row=1, column=col_idx)
            cell.font = self.styles['header_font']
            cell.fill = self.styles['header_fill']
            cell.alignment = self.styles['header_alignment']
            cell.border = self.styles['thin_border']

        # 添加数据行
        for _, row in df.iterrows():
            worksheet.append(row.tolist())

        # 应用数据样式
        for row_idx, row in enumerate(worksheet.iter_rows(min_row=2, max_row=len(df) + 1), 2):
            for col_idx, cell in enumerate(row, 1):
                cell.font = self.styles['data_font']
                cell.alignment = self.styles['data_alignment']
                cell.border = self.styles['thin_border']

                # 特殊格式处理
                header_name = headers[col_idx - 1]

                # 日期格式
                if header_name == '日期':
                    cell.number_format = self.styles['date_format']

                # 数字格式
                elif header_name in ['收盘价', '成交量(万手)']:
                    cell.number_format = self.styles['number_format']

                # 百分比格式
                elif header_name == '涨跌幅(%)':
                    cell.number_format = self.styles['percent_format']
                    # 根据正负设置背景色
                    if cell.value > 0:
                        cell.fill = self.styles['positive_fill']
                    elif cell.value < 0:
                        cell.fill = self.styles['negative_fill']

        # 自动调整列宽
        self.auto_adjust_columns(worksheet)

        # 冻结首行
        worksheet.freeze_panes = 'A2'

        # 添加筛选
        worksheet.auto_filter.ref = f"A1:{get_column_letter(worksheet.max_column)}{worksheet.max_row}"

    def auto_adjust_columns(self, worksheet: openpyxl.worksheet.worksheet.Worksheet) -> None:
        """
        自动调整列宽
        :param worksheet: 工作表对象
        """
        for col in worksheet.columns:
            max_length = 0
            column_letter = get_column_letter(col[0].column)

            for cell in col:
                try:
                    # 计算单元格内容的长度
                    cell_length = len(str(cell.value))
                    if cell_length > max_length:
                        max_length = cell_length
                except:
                    pass

            # 设置列宽，最大不超过50
            adjusted_width = min(max_length + 2, 50)
            worksheet.column_dimensions[column_letter].width = adjusted_width

    def export_stock_data(self, df: pd.DataFrame, output_file: str) -> str:
        """将股票数据导出到Excel文件，分为上海股票、深圳股票和北京股票三个工作表。

        Args:
            df (pd.DataFrame): 股票数据 DataFrame
            output_file (str): 输出文件路径

        Returns:
            str: 保存的文件路径
        """
        if df.empty:
            logging.warning("没有数据可以导出")
            return ""

        try:
            # 创建Excel工作簿
            wb = openpyxl.Workbook()

            # 删除默认的工作表
            if 'Sheet' in wb.sheetnames:
                del wb['Sheet']

            # 需要过滤掉的列
            columns_to_filter = ['涨跌额', '成交额', '最高', '今开', '昨收', '市盈率-动态', '最低', '涨速', '5分钟涨跌',
                                 '60日涨跌幅', '年初至今涨跌幅']

            # 首先过滤掉不需要的列
            filtered_df = df.drop(columns=[col for col in columns_to_filter if col in df.columns], errors='ignore')

            # 根据股票代码拆分数据
            sh_stocks = []  # 上海股票（代码以6、9、5开头）
            sz_stocks = []  # 深圳股票（代码以0、3开头）
            bj_stocks = []  # 北京股票（代码以4、8开头）

            for _, row in filtered_df.iterrows():
                stock_code = str(row['代码'])
                # 清理非数字字符
                clean_code = re.sub(r'\D', '', stock_code)

                # 根据股票代码的开头判断市场
                if clean_code.startswith(('6', '9', '5')):
                    sh_stocks.append(row)
                elif clean_code.startswith(('0', '3')):
                    sz_stocks.append(row)
                elif clean_code.startswith(('4', '8')):
                    bj_stocks.append(row)
                else:
                    # 未知市场的股票，默认放入深圳
                    sz_stocks.append(row)

            # 创建工作表并加入数据
            if sh_stocks:
                sh_df = pd.DataFrame(sh_stocks)
                ws_sh = wb.create_sheet(title="上海股票")
                self._add_data_to_sheet(ws_sh, sh_df, "上海股票")

            if sz_stocks:
                sz_df = pd.DataFrame(sz_stocks)
                ws_sz = wb.create_sheet(title="深圳股票")
                self._add_data_to_sheet(ws_sz, sz_df, "深圳股票")

            if bj_stocks:
                bj_df = pd.DataFrame(bj_stocks)
                ws_bj = wb.create_sheet(title="北京股票")
                self._add_data_to_sheet(ws_bj, bj_df, "北京股票")

            # 如果没有数据，删除所有工作表并添加一个空工作表
            if not wb.sheetnames:
                ws_empty = wb.create_sheet(title="无数据")
                ws_empty['A1'] = "没有符合条件的股票数据"

            # 保存文件
            wb.save(output_file)
            logging.info(f"数据已成功导出到 {output_file}，包含 {len(wb.sheetnames)} 个工作表")

            return output_file  # 返回保存路径

        except Exception as e:
            logging.error(f"导出Excel失败: {e}")
            return ""

    def _add_data_to_sheet(self, worksheet: openpyxl.worksheet.worksheet.Worksheet,
                           df: pd.DataFrame, title: str) -> None:
        """
        将DataFrame数据添加到工作表并应用样式

        Args:
            worksheet: 工作表对象
            df: 包含数据的DataFrame
            title: 工作表标题
        """
        # 添加标题行
        headers = list(df.columns)
        worksheet.append(headers)

        # 应用标题样式
        for col_idx, header in enumerate(headers, 1):
            cell = worksheet.cell(row=1, column=col_idx)
            cell.font = self.styles['header_font']
            cell.fill = self.styles['header_fill']
            cell.alignment = self.styles['header_alignment']
            cell.border = self.styles['thin_border']

        # 添加数据行
        for _, row in df.iterrows():
            worksheet.append(row.tolist())

        # 应用数据样式
        for row_idx, row in enumerate(worksheet.iter_rows(min_row=2, max_row=len(df) + 1), 2):
            for col_idx, cell in enumerate(row, 1):
                cell.font = self.styles['data_font']
                cell.border = self.styles['thin_border']

                # 特殊格式处理
                header_name = headers[col_idx - 1]

                # 日期格式
                if header_name == '日期':
                    cell.number_format = self.styles['date_format']

                # 数字格式
                elif header_name in ['收盘价', '成交量(万手)']:
                    cell.number_format = self.styles['number_format']
                    cell.alignment = self.styles['data_alignment']

                # 百分比格式
                elif header_name == '涨跌幅(%)':
                    cell.number_format = self.styles['percent_format']
                    cell.alignment = self.styles['data_alignment']
                    # 根据正负设置背景色
                    if cell.value > 0:
                        cell.fill = self.styles['positive_fill']
                    elif cell.value < 0:
                        cell.fill = self.styles['negative_fill']

        # 自动调整列宽
        self.auto_adjust_columns(worksheet)

        # 冻结首行
        worksheet.freeze_panes = 'A2'

        # 添加筛选
        worksheet.auto_filter.ref = f"A1:{get_column_letter(worksheet.max_column)}{worksheet.max_row}"


# 使用示例
if __name__ == "__main__":
    # 创建Excel导出器实例
    exporter = ExcelExporter()

    # 模拟信号数据
    signals_data: Dict[str, List[Dict[str, Any]]] = {
        'sz_signals': [
            {
                'stock_code': '000001',
                'stock_name': '平安银行',
                'date': '2025-06-04',
                'signal_name': '金叉信号',
                'description': '5日均线上穿20日均线',
                'rating': 'AA',
                'industry': '银行',
                'close_price': 15.32,
                'volume': 123.45,
                'change_pct': 0.025
            },
            {
                'stock_code': '000001',
                'stock_name': '平安银行',
                'date': '2025-06-04',
                'signal_name': '放量上涨',
                'description': '成交量放大超过5日均线2倍',
                'rating': 'AA',
                'industry': '银行',
                'close_price': 15.32,
                'volume': 123.45,
                'change_pct': 0.025
            },
            {
                'stock_code': '000002',
                'stock_name': '万科A',
                'date': '2025-06-04',
                'signal_name': '突破信号',
                'description': '股价突破近期高点',
                'rating': 'A+',
                'industry': '地产',
                'close_price': 18.75,
                'volume': 89.12,
                'change_pct': 0.018
            }
        ],
        'sh_signals': [
            {
                'stock_code': '600519',
                'stock_name': '贵州茅台',
                'date': '2025-06-04',
                'signal_name': '量价齐升',
                'description': '价格上涨同时成交量放大',
                'rating': 'AAA',
                'industry': '白酒',
                'close_price': 1800.50,
                'volume': 5.67,
                'change_pct': 0.012
            },
            {
                'stock_code': '601318',
                'stock_name': '中国平安',
                'date': '2025-06-04',
                'signal_name': '均线多头排列',
                'description': '5日>10日>20日>60日均线',
                'rating': 'AA+',
                'industry': '保险',
                'close_price': 62.45,
                'volume': 98.76,
                'change_pct': 0.015
            }
        ]
    }

    # 生成信号报告
    target_date = '2025-06-04'
    report_path = exporter.export_signals_report(signals_data, target_date)
    print(f"信号报告已生成: {report_path}")

    # 导出普通股票数据示例
    stock_data = pd.DataFrame({
        '日期': ['2025-06-01', '2025-06-02', '2025-06-03'],
        '股票代码': ['000001', '000002', '600519'],
        '股票名称': ['平安银行', '万科A', '贵州茅台'],
        '收盘价': [15.20, 18.50, 1795.00],
        '涨跌幅': [0.012, 0.008, 0.005]
    })

    exporter.export_stock_data(stock_data, "普通股票数据.xlsx")